package top.infopub.security.util;


import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

import org.apache.commons.lang3.StringUtils;


/**
 * XSS
 * @author Awoke
 * @version 2018年9月18日
 * @see XssHttpServletRequestWrapper
 * @since
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private Map<Object, Object> parameterMap = new HashMap<Object, Object>();

    private static final String END_WITH_JSON = "_json";

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        init();
    }

    @SuppressWarnings({"rawtypes"})
    private void init() {
        HttpServletRequest request = (HttpServletRequest)getRequest();
        if (request != null) {
            Map<String, String[]> map = request.getParameterMap();
            for (Entry entry : map.entrySet()) {
                // 参数k-v
                Object key = entry.getKey();
                Object value = entry.getValue();

                if (key != null && StringUtils.isNotBlank(key.toString())) {
                    if (value != null) {
                        // key是_json结尾  不转换value中的特殊字符
                        if (key.toString().endsWith(END_WITH_JSON)) {
                            parameterMap.put(key, value);
                        }
                        else {
                            if (value instanceof String[]) {
                                int length = ((String[])value).length;
                                String[] strs = new String[length];
                                for (int i = 0; i < length; i++ ) {
                                    strs[i] = HtmlClean.htmlSecurityEscape(((String[])value)[i]);
                                }
                                parameterMap.put(key, strs);
                            }
                            else if (value instanceof String) {
                                String temp = HtmlClean.htmlSecurityEscape((String)value);
                                parameterMap.put(key, temp);
                            }
                            else {
                                String temp = HtmlClean.htmlSecurityEscape(value.toString());
                                parameterMap.put(key, temp);
                            }
                        }
                    }
                }
            }
        }
    }

    @Override
    public String getParameter(String name) {
        Object object = parameterMap.get(name);
        if (object != null) {
            if (object instanceof String[]) {
                String[] array = (String[])object;
                if (array.length > 0) {
                    return array[0];
                }
            }
            else if (object instanceof String) {
                return (String)object;
            }
            else {
                return object.toString();
            }
        }
        return this.getRequest().getParameter(name);
    }

    @SuppressWarnings({"rawtypes", "unchecked"})
    @Override
    public Enumeration getParameterNames() {
        return Collections.enumeration(parameterMap.keySet());
    }

    @Override
    public String[] getParameterValues(String name) {
        Object object = parameterMap.get(name);
        if (object != null) {
            if (object instanceof String[]) {
                return (String[])object;
            }
            else if (object instanceof String) {
                return new String[] {(String)object};
            }
            else {
                return new String[] {object.toString()};
            }
        }
        return null;
    }

    @SuppressWarnings({"rawtypes", "unchecked"})
    @Override
    public Map getParameterMap() {
        return parameterMap;
    }

}
